package org.nhindirect.stagent;

import java.io.ByteArrayOutputStream;
import java.io.InputStream;
import java.security.Security;
import java.security.cert.X509Certificate;
import java.util.Enumeration;

import javax.mail.internet.InternetAddress;
import javax.mail.internet.MimeMessage;
import javax.mail.internet.MimeMultipart;

import org.nhindirect.stagent.ProtocolException.ProtocolError;
import org.nhindirect.stagent.cert.ICertificateService;
import org.nhindirect.stagent.cert.ICertificateStore;
import org.nhindirect.stagent.cert.X509CertificateEx;
import org.nhindirect.stagent.trust.ITrustSettingsStore;
import org.nhindirect.stagent.trust.TrustEnforcementStatus;
import org.nhindirect.stagent.trust.TrustError;
import org.nhindirect.stagent.trust.TrustException;
import org.nhindirect.stagent.trust.impl.TrustModel;

public class NHINDAgent 
{
	static MimeMultipart lastMMPart = null;
	
	Cryptographer m_cryptographer;
    ICertificateService m_internalCertService;
    ICertificateStore m_externalCertStore;
    ITrustSettingsStore m_trustSettings;
    TrustModel m_trustModel;
    boolean m_encryptionEnabled = true;
    TrustEnforcementStatus m_minTrustRequirement;
    String m_domain;
    NHINDAgentEventListener m_listener = null;
    
    static
    {
		Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
    }
    
    public NHINDAgent(String domain, ICertificateService internalCerts, ICertificateStore externalCerts, ITrustSettingsStore trustSettings)
    {
    	this(domain, internalCerts, externalCerts, trustSettings, TrustModel.Default, Cryptographer.Default);
    }

    public NHINDAgent(String domain, ICertificateService internalCerts, ICertificateStore externalCerts, ITrustSettingsStore trustSettings, TrustModel trustModel, Cryptographer cryptographer)
    {
        if (domain == null || domain.length() == 0 || internalCerts == null || externalCerts == null || trustSettings == null || trustModel == null || cryptographer == null)
        {
            throw new IllegalArgumentException();
        }

        this.m_domain = domain;
        this.m_internalCertService = internalCerts;
        this.m_externalCertStore = externalCerts;
        this.m_cryptographer = cryptographer;
        this.m_trustSettings = trustSettings;
        this.m_trustModel = trustModel;
        this.m_minTrustRequirement = TrustEnforcementStatus.Success_Offline;
    }

    public String getDomain()
    {
        return this.m_domain;
    }

    public Cryptographer getCryptographer()
    {
        return this.m_cryptographer;
    }

    public boolean isEncryptMessages()
    {

        return this.m_encryptionEnabled;
    }
     
    public void setEncryptMessages(boolean value)
    {
        this.m_encryptionEnabled = value;
    }


    public ICertificateStore getExternalCertStore()
    {
        return this.m_externalCertStore;
    }

    public ICertificateService getInternalCertService()
    {
        return this.m_internalCertService;
    }

    public ITrustSettingsStore getTrustSettings()
    {
        return this.m_trustSettings;
    }

    /// <summary>
    /// Messages must satisfy this minimum trust status
    /// </summary>
    public TrustEnforcementStatus getMinTrustRequirement()
    {
    	return this.m_minTrustRequirement;
    }
        
    
    public void setMinTrustRequirement(TrustEnforcementStatus value)
    {
        if (value.compareTo(TrustEnforcementStatus.Success_Offline) < 0)
        {
            throw new IllegalArgumentException();
        }
        this.m_minTrustRequirement = value;
    }

    public void setEventListener(NHINDAgentEventListener listener)
    {
    	m_listener = listener;
    }

    //-------------------------------------------------------------------
    //
    // INCOMING MESSAGE
    //
    //-------------------------------------------------------------------

    public String processIncoming(MimeMessage msg)
    {
		MimeEntity theEntity = new MimeEntity();
    	
		try
		{
			byte buffer[] = new byte[2048];
			InputStream inStream = msg.getRawInputStream();
			ByteArrayOutputStream oStream = new ByteArrayOutputStream();
			int count;
			while ((count = inStream.read(buffer)) > -1)
				oStream.write(buffer, 0, count);
			
    	    Enumeration<javax.mail.Header> en = msg.getAllHeaders();
    	    HeaderCollection headers = new HeaderCollection();
    	    while (en.hasMoreElements())
    	    {
    	    	javax.mail.Header hd = en.nextElement();
    	    	headers.add(new Header(hd.getName(), hd.getValue()));	        	    	
    	    }
    	    theEntity.setHeaders(headers);    	   
    	    
    	    theEntity.setBody(new Body(new String(oStream.toByteArray(), "ASCII")));	

		}
		catch (Exception e)
		{
			throw new NHINDException("Error converting to IncomingMessage", e);
		}
		
        return processIncoming(new IncomingMessage(theEntity));
    }    
    
    public String processIncoming(String messageText)
    {
        if (messageText == null || messageText.length() == 0)
        {
            throw new IllegalArgumentException();
        }

        MimeEntity entity = EntitySerializer.Default.deserialize(messageText);
        return processIncoming(new IncomingMessage(entity));
    }
    
    /// <summary>
    /// Incoming message from an external sender to somebody in this org...
    ///  - Decrypt, Verify Signature, etc
    /// </summary>
    public String processIncoming(IncomingMessage message)
    {          
    	if (message == null)
        {
            throw new IllegalArgumentException();
        }    	

        try
        {
            message.setAgent(this);
            message.validate();
            
            if (m_listener != null)
            	m_listener.preProcessIncoming(message);            

            this.processMessage(message);

            if (m_listener != null)
            	m_listener.postProcessIncoming(message);                

            message.updateTo();
            return EntitySerializer.Default.serialize(message);
        }
        catch (Exception error)
        {        	
        	NHINDException throwError = new NHINDException(error);
        	
            if (m_listener != null)
            	m_listener.errorIncoming(message, error);  
            throw throwError;  // rethrow error
        }
    }

    void processMessage(IncomingMessage message)
    {
        if (message.getSender() == null)
        {
            throw new TrustException(TrustError.UntrustedSender);
        }

        message.categorizeRecipients(this.getDomain());
        if (!message.hasDomainRecipients())
        {
        	throw new TrustException(TrustError.NoTrustedRecipients);
        }
        //
        // Map each address to its certificates/trust settings
        //
        this.bindAddresses(message);
        //
        // Extract signed content from the message
        //
        SignedEntity signedEntity = this.decryptSignedContent(message);
        
        
        //
        // Extract the signature. 
        //
        message.setSignature(this.m_cryptographer.deserializeSignatureEnvelope(signedEntity));
        //
        // Alter body to contain actual content. Also clean up mime headers on the message that were there to support
        // signatures etc
        //
        message.setHeaders(message.getHeaders().selectNonMimeHeaders());
        message.applyBody(signedEntity.getContent()); // this will merge in content + content specific mime headers
        //
        // Enforce trust requirements, including checking signatures
        //
        this.m_trustModel.enforce(message);
        //
        // Remove any untrusted recipients...
        //
        if (message.hasDomainRecipients())
        {
            message.categorizeRecipients(this.m_minTrustRequirement);
        }
        if (!message.hasDomainRecipients())
        {
            throw new TrustException(TrustError.NoTrustedRecipients);
        }
    }

    void bindAddresses(IncomingMessage message)
    {
        //
        // Retrieving the sender's certificate is optional
        //
        message.getSender().setCertificate(this.resolvePublicCert(message.getSender(), false));
        //
        // Bind each recpient's certs and trust settings
        //

        for (NHINDAddress recipient : message.getDomainRecipients())
        {
            recipient.setCertificate(this.resolvePrivateCert(recipient, false));
            recipient.setTrustAnchors(this.m_trustSettings.getTrustAnchorsIncoming(recipient));
        }
    }

    SignedEntity decryptSignedContent(IncomingMessage message)
    {
        SignedEntity signedEntity;
        if (this.m_encryptionEnabled)
        {
            //            
            // Decrypt and parse message body into a signature entity - the envelope that contains our data + signature
            // We can use the cert of any ONE of the recipients to decrypt
            //
            X509Certificate decryptingCert = message.getDomainRecipients().getFirstCertificate();
            if (decryptingCert == null)
            {
                throw new TrustException(TrustError.MissingRecipientCertificate);
            }
            MimeEntity decryptedEntity = this.m_cryptographer.decrypt(message, (X509CertificateEx)decryptingCert);
            //
            // Extract the signature envelope. That contains both the signature and the actual message content
            //
            signedEntity = SignedEntity.load(decryptedEntity);
        }
        else
        {
            signedEntity = SignedEntity.load(message);
        }

        return signedEntity;
    }

    //-------------------------------------------------------------------
    //
    // OUTGOING MESSAGE
    //
    //-------------------------------------------------------------------

    /// <summary>
    /// Outgoing message from somebody in this org to the outside world
    ///  - Sign, Encrypt
    /// </summary>
    public String processOutgoing(String messageText)
    {
        if (messageText == null || messageText.length() == 0)
        {
            throw new IllegalArgumentException();
        }

        MimeEntity entity = EntitySerializer.Default.deserialize(messageText);
        OutgoingMessage message = new OutgoingMessage(entity);
        
        message.validate();

        try
        {
            message.setAgent(this);


            if (m_listener != null)
            	m_listener.preProcessOutgoing(message);                

            this.processMessage(message);

            if (m_listener != null)
            	m_listener.postProcessOutgoing(message);     

            message.updateTo();
            return EntitySerializer.Default.serialize(message);
        }
        catch (Exception error)
        {        	        	
        	error.printStackTrace();
        	
        	NHINDException throwError = new NHINDException(error);
        	
            if (m_listener != null)
            	m_listener.errorOutgoing(message, error);  
            throw throwError;  // rethrow error
        }
    }

    void processMessage(OutgoingMessage message)
    {
        if (message.getSender() == null)
        {
            throw new ProtocolException(ProtocolError.MissingFrom);
        }

        this.bindAddresses(message);
        if (!message.hasRecipients())
        {
            throw new ProtocolException(ProtocolError.MissingTo);
        }
        //
        // Enforce the trust model.
        //
        this.m_trustModel.enforce(message);
        //
        // Remove any non-trusted recipients
        //
        message.categorizeRecipients(this.m_minTrustRequirement);
        if (!message.hasRecipients())
        {
            throw new TrustException(TrustError.NoTrustedRecipients);
        }
        //
        // Finally, sign and encrypt the message
        //
        this.signAndEncryptMessage(message);
    }

    void bindAddresses(OutgoingMessage message)
    {
        //
        // Retrieving the sender's private certificate is requied for encryption
        //
        message.getSender().setCertificate(this.resolvePrivateCert(message.getSender(), true));
        message.getSender().setTrustAnchors(this.m_trustSettings.getTrustAnchorsOutgoing(message.getSender()));
        //
        // Bind each recipient's certs
        //
        NHINDAddressCollection recipients = message.getRecipients();
        for (int i = 0, count = recipients.size(); i < count; ++i)
        {
            NHINDAddress recipient = recipients.get(i);
            recipient.setCertificate(this.resolvePublicCert(recipient, false));
        }
    }

    //
    // First sign, THEN encrypt the message
    //
    void signAndEncryptMessage(OutgoingMessage message)
    {
        SignedEntity signedEntity = this.m_cryptographer.sign(message, (X509CertificateEx)message.getSender().getCertificate());
                        
        if (this.m_encryptionEnabled)
        {
            MimeEntity encryptedEntity = this.m_cryptographer.encrypt(signedEntity, message.getRecipients().getCertificates());
            //
            // Alter message content to contain encrypted data
            //
            message.applyBody(encryptedEntity);
        }
        else
        {
            message.applyBody(signedEntity);
        }
    }

    InternetAddress resolveFrom(Message message)
    {
    	InternetAddress from = message.getParsedFrom();
        if (from == null)
        {
            throw new ProtocolException(ProtocolError.MissingFrom);
        }

        return from;
    }

    X509CertificateEx resolvePrivateCert(InternetAddress address, boolean required)
    {
        X509CertificateEx cert = null;
        try
        {
            cert = this.m_internalCertService.getPrivateCertificate(address);
            if (cert == null && required)
            {
                throw new TrustException(TrustError.UnknownRecipient);
            }
        }
        catch (Exception ex)
        {
            if (required)
            {
                if (this.m_listener != null)
                	m_listener.error(ex);
                
                throw new NHINDException(ex);
            }
        }

        return cert;
    }

    X509Certificate resolvePublicCert(InternetAddress address, boolean required) throws NHINDException
    {
        X509Certificate cert = null;
        try
        {
            if (address.getAddress().contains(this.m_domain))
            {
                //
                // Internal address
                //
                cert = this.m_internalCertService.getCertificate(address);
            }
            else
            {
                cert = this.m_externalCertStore.getCertificate(address);
            }

            if (cert == null && required)
            {
                throw new TrustException(TrustError.UnknownRecipient);
            }
        }
        catch (Exception ex)
        {
            if (required)
            {
                if (this.m_listener != null);
                	m_listener.error(ex);
                
                throw new NHINDException(ex);
            }

        }

        return cert;
    }

}
